package cn.sciento.websocket.registry;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import org.apache.commons.collections4.CollectionUtils;
import org.springframework.web.socket.WebSocketSession;

public class SessionRegistry {
    private static Map<String, WebSocketSession> sessionMap = new ConcurrentHashMap<>();

    private static Map<String, Long> userMap = new ConcurrentHashMap<>();

    private static String brokerId;

    public static void addSession(WebSocketSession session, String sessionId, Long userId) {
        sessionMap.put(sessionId, session);
        userMap.put(sessionId, userId);
    }

    public static void removeSession(String sessionId) {
        sessionMap.remove(sessionId);
        userMap.remove(sessionId);
    }

    public static List<WebSocketSession> getSession(List<String> sessionIds) {
        if (CollectionUtils.isEmpty(sessionIds)) {
            return new ArrayList<>();
        }
        List<WebSocketSession> result = new ArrayList<>();
        sessionIds.forEach(sessionId -> result.add(sessionMap.get(sessionId)));
        return result;
    }

    public static WebSocketSession getSession(String sessionId) {
        return sessionMap.get(sessionId);
    }

    public static Long getUser(String sessionId) {
        return userMap.get(sessionId);
    }

    public static String getBrokerId() {
        return brokerId;
    }

    public static void setBrokerId(String brokerId) {
        SessionRegistry.brokerId = brokerId;
    }
}
